#!/usr/bin/env python3
"""Transform utilities for point cloud processing."""

import numpy as np
import random
import json
import os
from pathlib import Path
from loadJson import _save_neutral_gripper,_load_gripper_transform, _load_finger_length,_load_emb_name_map, _get_neutral_gripper
from scipy.spatial.transform import Rotation as R
from sklearn.cluster import KMeans
import time
from time_it import timeit
import transforms3d as t3d
from typing import Tuple
time_now=int(time.time() * 1000) %1000   
class Transform:
    current_biases = {}  # 存储当前手部的 z 偏移量
    transform_file=""
    stick_bias_vec=np.array([0,0,0])
    state_action_vec_on_collect_written=False
    @staticmethod
    def stick_faces(points: np.ndarray, norm_vecs: list, bbox_local_colored: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Move transformed points so specified faces match original bbox faces."""
        if len(norm_vecs) == 1: 
            target_point, current_point = Transform._get_face_center(bbox_local_colored, norm_vecs[0]), Transform._get_face_center(points, norm_vecs[0])
        elif len(norm_vecs) == 2: 
            target_point, current_point = Transform._get_intersection_midpoint(bbox_local_colored, norm_vecs), Transform._get_intersection_midpoint(points, norm_vecs)
        elif len(norm_vecs) == 3: 
            target_point, current_point = Transform._get_corner_point(bbox_local_colored, norm_vecs), Transform._get_corner_point(points, norm_vecs)
        else: 
            return points
        points[:, :3] += target_point - current_point
        return points, target_point - current_point
    
    @staticmethod
    def _get_face_center(points: np.ndarray, norm_vec: np.ndarray) -> np.ndarray:
        """Get center of face specified by normal vector."""
        min_coords, max_coords = np.min(points[:, :3], axis=0), np.max(points[:, :3], axis=0)
        face_center = np.zeros(3)
        for i in range(3):
            if norm_vec[i] > 0.5: face_center[i] = min_coords[i]
            elif norm_vec[i] < -0.5: face_center[i] = max_coords[i]
            else: face_center[i] = (min_coords[i] + max_coords[i]) / 2
        return face_center
    
    @staticmethod
    def _get_intersection_midpoint(points: np.ndarray, norm_vecs: list) -> np.ndarray:
        """Get midpoint of intersection line of two faces."""
        face1_center = Transform._get_face_center(points, norm_vecs[0])
        face2_center = Transform._get_face_center(points, norm_vecs[1])
        return (face1_center + face2_center) / 2
    
    @staticmethod
    def _get_corner_point(points: np.ndarray, norm_vecs: list) -> np.ndarray:
        """Get corner point where three faces meet."""
        min_coords, max_coords = np.min(points[:, :3], axis=0), np.max(points[:, :3], axis=0)
        corner = np.zeros(3)
        for i in range(3):
            if any(norm_vec[i] > 0.5 for norm_vec in norm_vecs): corner[i] = max_coords[i]
            elif any(norm_vec[i] < -0.5 for norm_vec in norm_vecs): corner[i] = min_coords[i]
            else: corner[i] = (min_coords[i] + max_coords[i]) / 2
        return corner
    
    @staticmethod
    @timeit()
    def get_transforms(part_points_local: np.ndarray, norm_vecs: list, bbox_local_colored: np.ndarray, part_name: str,
                       stick_faces: bool = False, z_boundary: bool = False, transform_file: str | None = None,
                       neutral_stretch: str | None = None, want_neutral_stretch: bool = False,
                       episode: int | None = None, ts: int | None = None,emb_char = ''):
        """Apply random transforms based on JSON configuration."""
        applied_tools = set()
        if want_neutral_stretch:
            neutral_data = {}
            for key in ['hand', 'finger']:
                if part_name.endswith(key):
                    bbox_min, bbox_max = np.min(bbox_local_colored[:, :3], axis=0), np.max(bbox_local_colored[:, :3], axis=0)
                    stretch_dims = bbox_max - bbox_min
                    pyramid_areas = []
                    for norm_vec in [[1,0,0], [-1,0,0], [0,1,0], [0,-1,0], [0,0,1], [0,0,-1]]:
                        aligned_bbox = Transform.align_norm_to_z(bbox_local_colored, np.array(norm_vec))
                        aligned_min, aligned_max = np.min(aligned_bbox[:, :3], axis=0), np.max(aligned_bbox[:, :3], axis=0)
                        area = (aligned_max[0] - aligned_min[0]) * (aligned_max[1] - aligned_min[1])
                        pyramid_areas.append(area)
                    neutral_data[key] = {"stretch": stretch_dims.tolist(), "pyramid": pyramid_areas}
            _save_neutral_gripper(neutral_data, emb_char)

        # If no transform file is specified, skip transforming
        if not transform_file:
            return part_points_local

        # When a transform file is provided, neutral_stretch must be provided as well
        assert neutral_stretch is not None and len(str(neutral_stretch)) > 0, \
            f"neutral_stretch must be provided when transform_file is specified (transform_file={transform_file})"

        Transform.force_sync_stretch(transform_file, neutral_stretch)
        config = _load_gripper_transform(transform_file)
        for key in config.keys():
            if part_name.endswith(key):  # contains match
                # Defer stretch to the end to ensure absolute target length after all other transforms
                deferred_stretch_params = None
                for transform_config in config[key]:
                    tool_name = transform_config['tool']
                    if tool_name in applied_tools or tool_name in ['flyPoint']:
                        continue
                    param_range = transform_config['param_range']
                    if hasattr(Transform, tool_name):
                        tool_func = getattr(Transform, tool_name)
                        params = []
                        for param_name, param_value in param_range.items():
                            if isinstance(param_value, list) and len(param_value) >= 2 and tool_name != 'color':
                                if len(param_value) == 2:
                                    seed_flag = 3
                                elif len(param_value) == 3:
                                    seed_flag = param_value[2]
                                    param_value = param_value[:2]
                                Transform._set_random_seed(seed_flag, episode, ts,f"{key},{transform_config['tool']},{param_name}")
                                param_value = random.uniform(*param_value)
                            params.append(param_value)
                        
                        # Convert absolute meters to relative scale for stretch and pyramid
                        if tool_name == 'stretch':
                            # Use current bbox of the part (after prior transforms in this loop)
                            cur_min, cur_max = np.min(part_points_local[:, :3], axis=0), np.max(part_points_local[:, :3], axis=0)
                            cur_dims = cur_max - cur_min
                            # print("part name:", part_name)
                            # print(f"bbox_min: {cur_min}, bbox_max: {cur_max}")
                            # print(f"params: {params}")
                            # print(f"orig_dims: {cur_dims}")
                            # For force_sync_stretch, use absolute target dimensions directly
                            if len(param_range['x']) == 3 and param_range['x'][2] == 1:
                                deferred_stretch_params = [params[0], params[1], params[2], True]
                            else:
                                deferred_stretch_params = [params[0]/cur_dims[0], params[1]/cur_dims[1], params[2]/cur_dims[2], False]
                            # Defer applying stretch until after other tools
                            continue
                        elif tool_name == 'pyramid':
                            new_params = []
                            for i, norm_vec in enumerate([[1,0,0], [-1,0,0], [0,1,0], [0,-1,0], [0,0,1], [0,0,-1]]):
                                aligned_bbox = Transform.align_norm_to_z(bbox_local_colored, np.array(norm_vec))
                                aligned_min, aligned_max = np.min(aligned_bbox[:, :3], axis=0), np.max(aligned_bbox[:, :3], axis=0)
                                orig_area = (aligned_max[0] - aligned_min[0]) * (aligned_max[1] - aligned_min[1])
                                new_params.append((params[i]/orig_area)**0.5)
                            params = new_params
                        
                        part_points_local = tool_func(part_points_local, *params)
                        applied_tools.add(tool_name)
                    else:
                        print(f"Warning: Tool {tool_name} not found in Transform class")
                # Apply stretch last if it was present
                if deferred_stretch_params is not None:
                    part_points_local = Transform.stretch(part_points_local, *deferred_stretch_params)
                    applied_tools.add('stretch')
        if stick_faces:
            part_points_local, Transform.stick_bias_vec = Transform.stick_faces(part_points_local, norm_vecs, bbox_local_colored)
        if z_boundary:
            part_points_local = Transform.boundary(part_points_local, part_name,bbox_local_colored,episode)
        
        for key in config.keys():
            if part_name.endswith(key):
                for transform_config in config[key]:
                    if transform_config['tool'] == 'flyPoint':
                        param_range = transform_config['param_range']
                        part_points_local = Transform.flyPoint(part_points_local, **param_range)
                        break
        
        return part_points_local
    
    @staticmethod
    def scale(points: np.ndarray, scale: float):
        """Scale points around their center (bounding-box center)."""
        scaled_points = points.copy()
        min_coords = np.min(points[:, :3], axis=0)
        max_coords = np.max(points[:, :3], axis=0)
        center = (min_coords + max_coords) / 2.0
        scaled_points[:, :3] = center + (points[:, :3] - center) * scale
        return scaled_points
    
    @staticmethod
    def rpy(points: np.ndarray, r: float, p: float, y: float):
        """Rotate points by roll, pitch, yaw angles."""
        from scipy.spatial.transform import Rotation as R
        rotation = R.from_euler('xyz', [r, p, y])
        points[:, :3] = rotation.apply(points[:, :3])
        return points
    
    @staticmethod
    def flyPoint(points: np.ndarray, percentage: float, sigma: float):
        """Randomly deviate a percentage of points by sigma deviation."""
        mask = np.random.random(len(points)) < float(percentage)
        points[mask, :3] += np.random.normal(0, float(sigma), (np.sum(mask), 3))
        return points
    
    @staticmethod
    def stretch(points: np.ndarray, x: float, y: float, z: float, use_absolute: bool = False):
        """Stretch points in x, y, z directions."""
        # print(f"stretch x:{x:.3f} y:{y:.3f} z:{z:.3f} {'(absolute)' if use_absolute else '(relative)'}")
        if use_absolute:
            # Use absolute target dimensions
            min_coords = np.min(points[:, :3], axis=0)
            max_coords = np.max(points[:, :3], axis=0)
            cur_dims = max_coords - min_coords
            center = (min_coords + max_coords) / 2.0
            scale_x = x / cur_dims[0] if cur_dims[0] > 0 else 1.0
            scale_y = y / cur_dims[1] if cur_dims[1] > 0 else 1.0
            scale_z = z / cur_dims[2] if cur_dims[2] > 0 else 1.0
            points[:, 0] = center[0] + (points[:, 0] - center[0]) * scale_x
            points[:, 1] = center[1] + (points[:, 1] - center[1]) * scale_y
            points[:, 2] = center[2] + (points[:, 2] - center[2]) * scale_z
        else:
            # Scale around the point cloud's center (bounding-box center) to avoid shifting
            min_coords = np.min(points[:, :3], axis=0)
            max_coords = np.max(points[:, :3], axis=0)
            center = (min_coords + max_coords) / 2.0
            points[:, 0] = center[0] + (points[:, 0] - center[0]) * x
            points[:, 1] = center[1] + (points[:, 1] - center[1]) * y
            points[:, 2] = center[2] + (points[:, 2] - center[2]) * z
        bbox_min, bbox_max = np.min(points[:, :3], axis=0), np.max(points[:, :3], axis=0)
        print(f"after stretch diff: {(bbox_max-bbox_min)[0]:.3f}",end="\r",flush=True)
        return points
    
    @staticmethod
    def pyramid(points: np.ndarray, xp: float, xm: float, yp: float, ym: float, zp: float, zm: float):
        """Apply pyramid transformation to stretch faces horizontally based on scale factors."""
        vecmap = {
            tuple([1,0,0]): xp,
            tuple([-1,0,0]): xm,
            tuple([0,1,0]): yp,
            tuple([0,-1,0]): ym,
            tuple([0,0,1]): zp,
            tuple([0,0,-1]): zm,
        }
        
        for norm_vec, scale in vecmap.items():
            if scale != 1.0:
                points = Transform._apply_pyramid_transform(points, np.array(norm_vec), scale)
        return points
    
    @staticmethod
    def _apply_pyramid_transform(points: np.ndarray, norm_vec: np.ndarray, scale: float) -> np.ndarray:
        """Apply pyramid transformation for a specific face direction."""
        aligned_points = Transform.align_norm_to_z(points.copy(), norm_vec)
        min_coords = np.min(aligned_points[:, :3], axis=0)
        max_coords = np.max(aligned_points[:, :3], axis=0)
        z_min, z_max = min_coords[2], max_coords[2]
        x_center = (min_coords[0] + max_coords[0]) / 2
        y_center = (min_coords[1] + max_coords[1]) / 2
        for i in range(len(aligned_points)):
            z_cur = aligned_points[i, 2]
            if z_max != z_min:
                z_ratio = (z_cur - z_min) / (z_max - z_min)
                current_scale = 1.0 + (scale - 1.0) * z_ratio
            else:
                current_scale = scale
            aligned_points[i, 0] = x_center + (aligned_points[i, 0] - x_center) * current_scale
            aligned_points[i, 1] = y_center + (aligned_points[i, 1] - y_center) * current_scale
        return Transform.align_norm_to_z(aligned_points, norm_vec, inv=True)
    
   
    
    @staticmethod
    def align_norm_to_z(points: np.ndarray, norm_vec: np.ndarray, inv: bool = False) -> np.ndarray:
        points=points.copy()
        """Transform points so normal vector points to z+, maintaining relative positions."""
        if not inv: rotation = R.align_vectors([[0, 0, 1]], [norm_vec])[0]
        else: rotation = R.align_vectors([norm_vec], [[0, 0, 1]])[0]
        points[:, :3] = rotation.apply(points[:, :3])
        return points
    @staticmethod
    def global_points_to_local(points: np.ndarray,rpy,center) -> np.ndarray:
        """Transform points from global to local coordinate system using key_pose RPY as Z+ axis. Support [N, 3] or [N, 6]"""
        rotation = R.from_euler('xyz', rpy)
        local_positions = rotation.inv().apply(points[:, :3] - center)
        return np.column_stack([local_positions, points[:, 3:]])

    @staticmethod
    def local_points_to_global(local_points: np.ndarray,rpy,center) -> np.ndarray:
        """Convert local points back to global coordinate system. Support [N, 3] or [N, 6]"""
        rotation = R.from_euler('xyz', rpy)
        global_positions = rotation.apply(local_points[:, :3]) + center
        return np.column_stack([global_positions, local_points[:, 3:]])
    @staticmethod
    def get_existing_color_param(points: np.ndarray, num_colors: int = 3):
        """Analyze existing color distribution and return mainColors, weights, sigma for current points."""
        colors = points[:, 3:6]
        kmeans = KMeans(n_clusters=num_colors, random_state=42)
        cluster_labels = kmeans.fit_predict(colors)
        cluster_centers = kmeans.cluster_centers_
        weights = np.bincount(cluster_labels) / len(points)
        sigma = []
        for i in range(num_colors):
            cluster_points = colors[cluster_labels == i]
            if len(cluster_points) > 0:
                sigma.append(np.std(cluster_points, axis=0).mean())
            else:
                sigma.append(0.0)
        mainColors_list = [list(np.round(center, 2)) for center in cluster_centers]
        weights_list = list(np.round(weights, 2))
        sigma_list = list(np.round(sigma, 2))
        # print(f"Current part mainColors: {mainColors_list}, weights: {weights_list}, sigma: {sigma_list}")
        return mainColors_list, weights_list, sigma_list
    
    @staticmethod
    def color(points: np.ndarray, mainColors: list, weights: list, sigma: list):
        """Randomly recolor each point based on main colors, weights and deviation."""
        # Transform.get_existing_color_param(points)
        if len(mainColors) == 0 or len(weights) == 0 or len(sigma) == 0:
            return points
        if sum(weights) != 1:
            weights = [weight / sum(weights) for weight in weights]
        color_indices = np.random.choice(len(mainColors), size=len(points), p=weights)
        for i, color_idx in enumerate(color_indices):
            points[i, 3:6] = np.array(mainColors[color_idx]) + np.random.normal(0, sigma[color_idx], 3)
        return points 

    @staticmethod
    def move(points: np.ndarray, x: float, y: float, z: float):
        """Move points by given x, y, z offsets."""
        points[:, 0] += x
        points[:, 1] += y
        points[:, 2] += z
        return points
    
    @staticmethod
    def _set_random_seed(seed_flag: int, episode: int = None, ts: int = None,tag:str=None):
        """Set random seed based on flag value. seed_flag: 0=fixed(42), 1=episode, 2=ts, 3=episode*1000+ts"""
        assert seed_flag == 0 or (seed_flag == 1 and episode is not None) or (seed_flag == 2 and ts is not None) or (seed_flag == 3 and episode is not None and ts is not None), f"Invalid seed_flag {seed_flag} or missing required params: episode={episode}, ts={ts}"
        seed = 42 if seed_flag == 0 else (episode+1)*1000 if seed_flag == 1 else ts+1 if seed_flag == 2 else (episode+1) * 1000 + (ts+1)
        seed += (time_now*1000000)%(2**32-1)
        # print(f"seed: {seed} time_now: {time_now} episode: {episode} ts: {ts} tag: {tag}")
        random.seed(seed); np.random.seed(seed)

    @staticmethod
    def boundary(points: np.ndarray, part_name: str,bbox_local_colored: np.ndarray,episode:int,norm_vec:np.ndarray=np.array([-1,0,0])):
        """Adjust norm_vec face boundary for hand parts to maintain original bbox alignment."""
        points=Transform.align_norm_to_z(points,norm_vec)
        bbox_aligned=Transform.align_norm_to_z(bbox_local_colored,norm_vec)
        if episode not in Transform.current_biases:
            Transform.current_biases[episode]={}
        if part_name[0] not in Transform.current_biases[episode]:
            assert part_name[0] in ["l","r"],f"part_name must start with l or r, but got {part_name}"
            assert part_name.endswith('lfinger') or part_name.endswith('rfinger'),f"current_biases{episode} {part_name[0]} must be set by part_name ending with lfinger or rfinger, but got {part_name}" 
            Transform.current_biases[episode][part_name[0]] = bbox_aligned[:, 2].min() - points[:, 2].min()
        points=Transform.move(points,0,0,Transform.current_biases[episode][part_name[0]])
        points=Transform.align_norm_to_z(points,norm_vec,inv=True)
        return points

    @staticmethod
    def force_sync_stretch(transform_file: str, neutral_stretch: str):
        neutral_data = _get_neutral_gripper(neutral_stretch)
        transform_data = json.load(open(Path(__file__).parent / f'gripper_rand_transform_{transform_file}.json'))
        for part in ['hand', 'finger']:
            if part in neutral_data and part in transform_data:
                neutral_x = neutral_data[part]['stretch'][0]
                for tool_config in transform_data[part]:
                    if tool_config['tool'] in ['stretch','pyramid']:
                        for param_name, param_value in tool_config['param_range'].items():
                            if param_name == 'x' and tool_config['tool'] == 'stretch':
                                tool_config['param_range']['x'] = [neutral_x, neutral_x, 1]
                            elif isinstance(param_value, (int, float)):
                                tool_config['param_range'][param_name] = [param_value, param_value, 1]
                            elif isinstance(param_value, list) and len(param_value) == 2:
                                tool_config['param_range'][param_name] = [param_value[0], param_value[1], 1]
                            elif isinstance(param_value, list) and len(param_value) == 3 and param_value[2] != 1:
                                tool_config['param_range'][param_name] = [param_value[0], param_value[1], 1]
        json.dump(transform_data, open(Path(__file__).parent / f'gripper_rand_transform_{transform_file}.json', 'w'), indent=2)
    @staticmethod
    def write_link_finger_7D_on_collect(ee_linkend7D,ee_fingerend7D,emb_fullname):
        emb_dict=_load_finger_length()[emb_fullname]
        # if emb_dict.get("state_pose_to_end_vec",None) is not None and Transform.state_action_vec_on_collect_written:
        #     print(f"state_pose_to_end_vec already written {emb_dict['state_pose_to_end_vec']}, skip")
        #     return
        relative_vec_global =  np.array(ee_fingerend7D[:3]) -  np.array(ee_linkend7D[:3])
        gripper_quat = np.array(ee_linkend7D[3:])
        rotation = R.from_quat(gripper_quat)
        relative_vec_local = rotation.apply(relative_vec_global)
        emb_dict["state_pose_to_end_vec"] = relative_vec_local.tolist()
        print(f"global:{np.round(ee_fingerend7D[:3],3)}-{np.round(ee_linkend7D[:3],3)}={np.round(relative_vec_global,3)} local:{np.round(relative_vec_local,3)} quad {np.round(ee_linkend7D[3:],3)}")
        if input()=="":
            pass
        Transform.state_action_vec_on_collect_written=True
    @staticmethod
    def write_link_finger_16D_on_collect(ee_linkend16D,ee_fingerend16D,emb_fullname):
        assert len(ee_linkend16D)==16 and len(ee_fingerend16D)==16,f"ee_linkend16D and ee_fingerend16D must be 16D, but got {len(ee_linkend16D)} and {len(ee_fingerend16D)}"
        left_linkend=ee_linkend16D[:7].copy()
        right_linkend=ee_linkend16D[8:15].copy()
        left_fingerend=ee_fingerend16D[:7].copy()
        right_fingerend=ee_fingerend16D[8:15].copy()
        Transform.write_link_finger_7D_on_collect(left_linkend,left_fingerend,emb_fullname)
        # Transform.write_link_finger_7D_on_collect(right_linkend,right_fingerend,emb_fullname)
    @staticmethod
    def finger_to_link_7D_on_eval(ee_fingerend7D,emb_fullname):
        emb_dict=_load_finger_length()[emb_fullname]
        relative_vec=emb_dict["state_pose_to_end_vec"]
        finger_pos = np.array(ee_fingerend7D[:3])
        finger_quat = np.array(ee_fingerend7D[3:])
        rotation = R.from_quat(finger_quat)
        offset_global = rotation.apply(relative_vec)
        return np.concatenate([finger_pos + offset_global, finger_quat])
    @staticmethod
    def finger_to_link_16D_on_eval(ee_fingerend16D,emb_fullname):
        ee_fingerend16D=ee_fingerend16D.copy()
        ee_fingerend16D[:7]=Transform.finger_to_link_7D_on_eval(ee_fingerend16D[:7],emb_fullname)
        ee_fingerend16D[8:15]=Transform.finger_to_link_7D_on_eval(ee_fingerend16D[8:15],emb_fullname)
        return ee_fingerend16D



